#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import shutil
import psutil
import random
import asyncio
import subprocess
from datetime import datetime


class Util:

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def procCreate(cpuCores, *kargs, **kwargs):
        if cpuCores is not None:
            curProc = psutil.Process()
            originalAffinity = curProc.cpu_affinity()
            if cpuCores < len(originalAffinity):
                newAffinity = random.sample(originalAffinity, cpuCores)
                curProc.cpu_affinity(newAffinity)
            else:
                cpuCores = None
        try:
            return subprocess.Popen(*kargs, **kwargs)
        finally:
            if cpuCores is not None:
                curProc.cpu_affinity(originalAffinity)

    @staticmethod
    def procWait(proc, *kargs, **kwargs):
        proc.wait()
        if proc.returncode != 0:
            raise subprocess.CalledProcessError(proc.returncode, proc.args)

    @staticmethod
    def procCommunicate(proc, *kargs, **kwargs):
        stdout, stderr = proc.communicate()
        if proc.returncode != 0:
            raise subprocess.CalledProcessError(proc.returncode, proc.args, output=stdout, stderr=stderr)
        return (stdout, stderr)

    @staticmethod
    async def asyncProcCreateExec(cpuCores, *kargs, **kwargs):
        # FIXME: we have not protect affinity, there's race conditions to be solved
        if cpuCores is not None:
            curProc = psutil.Process()
            originalAffinity = curProc.cpu_affinity()
            if cpuCores < len(originalAffinity):
                newAffinity = random.sample(originalAffinity, cpuCores)
                curProc.cpu_affinity(newAffinity)
            else:
                cpuCores = None
        try:
            return await asyncio.create_subprocess_exec(*kargs, **kwargs)
        finally:
            if cpuCores is not None:
                curProc.cpu_affinity(originalAffinity)

    @staticmethod
    async def asyncProcCreateShell(cpuCores, *kargs, **kwargs):
        # FIXME: we have not protect affinity, there's race conditions to be solved
        if cpuCores is not None:
            curProc = psutil.Process()
            originalAffinity = curProc.cpu_affinity()
            if cpuCores < len(originalAffinity):
                newAffinity = random.sample(originalAffinity, cpuCores)
                curProc.cpu_affinity(newAffinity)
            else:
                cpuCores = None
        try:
            return await asyncio.create_subprocess_shell(*kargs, **kwargs)
        finally:
            if cpuCores is not None:
                curProc.cpu_affinity(originalAffinity)

    @staticmethod
    async def asyncProcWait(proc, program_args, *kargs, **kwargs):
        await proc.wait()
        if proc.returncode != 0:
            raise subprocess.CalledProcessError(proc.returncode, program_args)

    @staticmethod
    async def asyncProcCommunicate(proc, cmd, *kargs, **kwargs):
        stdout, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise subprocess.CalledProcessError(proc.returncode, cmd, output=stdout, stderr=stderr)
        return (stdout, stderr)


class WindowsFileTime:
    """
    Utility class to deal with Windows FILETIME structures.
    See: https://msdn.microsoft.com/en-us/library/ms724284(v=vs.85).aspx
    """

    ticks_per_seconds = 10000000  # 1 tick every 100 nanoseconds
    epoch_delta = 11644473600  # 3600 * 24 * ((1970 - 1601) * 365 + 89)

    def __init__(self, timestamp=None):
        self.timestamp = timestamp

    def __repr__(self):
        return "<{}>: {}".format(self.__class__.__name__, self.timestamp)

    @classmethod
    def from_hex(cls, hexvalue):
        timestamp = int(hexvalue, 16)
        return WindowsFileTime(timestamp)

    def to_hex(self):
        return "{:x}".format(self.timestamp)

    @classmethod
    def from_unix_timestamp(cls, timestamp):
        timestamp = timestamp + cls.epoch_delta
        timestamp = int(timestamp * cls.ticks_per_seconds)
        return WindowsFileTime(timestamp)

    def to_unix_timestamp(self):
        if not self.timestamp:
            raise ValueError("No timestamp set")
        unix_ts = self.timestamp / self.ticks_per_seconds
        unix_ts = unix_ts - self.epoch_delta
        return unix_ts

    def to_date_time(self):
        return datetime.fromtimestamp(self.to_unix_timestamp())
